

<!DOCTYPE html>


<html lang="en" data-content_root="" >

  <head>
    <meta charset="utf-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <title>brevitas.core.bit_width.parameter &#8212; Brevitas Documentation - v0.10.0</title>
  
  
  
  <script data-cfasync="false">
    document.documentElement.dataset.mode = localStorage.getItem("mode") || "";
    document.documentElement.dataset.theme = localStorage.getItem("theme") || "";
  </script>
  
  <!-- Loaded before other Sphinx assets -->
  <link href="../../../../_static/styles/theme.css?digest=3ee479438cf8b5e0d341" rel="stylesheet" />
<link href="../../../../_static/styles/bootstrap.css?digest=3ee479438cf8b5e0d341" rel="stylesheet" />
<link href="../../../../_static/styles/pydata-sphinx-theme.css?digest=3ee479438cf8b5e0d341" rel="stylesheet" />

  
  <link href="../../../../_static/vendor/fontawesome/6.5.2/css/all.min.css?digest=3ee479438cf8b5e0d341" rel="stylesheet" />
  <link rel="preload" as="font" type="font/woff2" crossorigin href="../../../../_static/vendor/fontawesome/6.5.2/webfonts/fa-solid-900.woff2" />
<link rel="preload" as="font" type="font/woff2" crossorigin href="../../../../_static/vendor/fontawesome/6.5.2/webfonts/fa-brands-400.woff2" />
<link rel="preload" as="font" type="font/woff2" crossorigin href="../../../../_static/vendor/fontawesome/6.5.2/webfonts/fa-regular-400.woff2" />

    <link rel="stylesheet" type="text/css" href="../../../../_static/pygments.css" />
    <link rel="stylesheet" type="text/css" href="../../../../_static/sg_gallery.css" />
  
  <!-- Pre-loaded scripts that we'll load fully later -->
  <link rel="preload" as="script" href="../../../../_static/scripts/bootstrap.js?digest=3ee479438cf8b5e0d341" />
<link rel="preload" as="script" href="../../../../_static/scripts/pydata-sphinx-theme.js?digest=3ee479438cf8b5e0d341" />
  <script src="../../../../_static/vendor/fontawesome/6.5.2/js/all.min.js?digest=3ee479438cf8b5e0d341"></script>

    <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
    <script src="../../../../_static/jquery.js"></script>
    <script src="../../../../_static/underscore.js"></script>
    <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
    <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/sphinx_highlight.js"></script>
    <script crossorigin="anonymous" integrity="sha256-Ae2Vz/4ePdIu6ZyI/5ZGsYnb+m0JlOmKPjt6XZ9JJkA=" src="https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.4/require.min.js"></script>
    <script>DOCUMENTATION_OPTIONS.pagename = '_modules/brevitas/core/bit_width/parameter';</script>
    <script>
        DOCUMENTATION_OPTIONS.theme_version = '0.15.3';
        DOCUMENTATION_OPTIONS.theme_switcher_json_url = 'https://xilinx.github.io/brevitas/dev/_static/versions.json';
        DOCUMENTATION_OPTIONS.theme_switcher_version_match = 'v0.10.0';
        DOCUMENTATION_OPTIONS.show_version_warning_banner = false;
        </script>
    <link rel="author" title="About these documents" href="../../../../about.html" />
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" />
  <meta name="viewport" content="width=device-width, initial-scale=1"/>
  <meta name="docsearch:language" content="en"/>
  </head>
  
  
  <body data-bs-spy="scroll" data-bs-target=".bd-toc-nav" data-offset="180" data-bs-root-margin="0px 0px -60%" data-default-mode="">

  
  
  <div id="pst-skip-link" class="skip-link d-print-none"><a href="#main-content">Skip to main content</a></div>
  
  <div id="pst-scroll-pixel-helper"></div>
  
  <button type="button" class="btn rounded-pill" id="pst-back-to-top">
    <i class="fa-solid fa-arrow-up"></i>Back to top</button>

  
  <input type="checkbox"
          class="sidebar-toggle"
          id="pst-primary-sidebar-checkbox"/>
  <label class="overlay overlay-primary" for="pst-primary-sidebar-checkbox"></label>
  
  <input type="checkbox"
          class="sidebar-toggle"
          id="pst-secondary-sidebar-checkbox"/>
  <label class="overlay overlay-secondary" for="pst-secondary-sidebar-checkbox"></label>
  
  <div class="search-button__wrapper">
    <div class="search-button__overlay"></div>
    <div class="search-button__search-container">
<form class="bd-search d-flex align-items-center"
      action="../../../../search.html"
      method="get">
  <i class="fa-solid fa-magnifying-glass"></i>
  <input type="search"
         class="form-control"
         name="q"
         id="search-input"
         placeholder="Search the docs ..."
         aria-label="Search the docs ..."
         autocomplete="off"
         autocorrect="off"
         autocapitalize="off"
         spellcheck="false"/>
  <span class="search-button__kbd-shortcut"><kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd>K</kbd></span>
</form></div>
  </div>

  <div class="pst-async-banner-revealer d-none">
  <aside id="bd-header-version-warning" class="d-none d-print-none" aria-label="Version warning"></aside>
</div>

  
    <header class="bd-header navbar navbar-expand-lg bd-navbar d-print-none">
<div class="bd-header__inner bd-page-width">
  <button class="sidebar-toggle primary-toggle" aria-label="Site navigation">
    <span class="fa-solid fa-bars"></span>
  </button>
  
  
  <div class="col-lg-3 navbar-header-items__start">
    
      <div class="navbar-item">

  

<a class="navbar-brand logo" href="../../../../index.html">
  
  
  
  
  
    
    
      
    
    
    <img src="../../../../_static/brevitas_logo_black.svg" class="logo__image only-light" alt="Brevitas Documentation - v0.10.0 - Home"/>
    <script>document.write(`<img src="../../../../_static/brevitas_logo_white.svg" class="logo__image only-dark" alt="Brevitas Documentation - v0.10.0 - Home"/>`);</script>
  
  
</a></div>
    
  </div>
  
  <div class="col-lg-9 navbar-header-items">
    
    <div class="me-auto navbar-header-items__center">
      
        <div class="navbar-item">
<nav class="navbar-nav">
  <ul class="bd-navbar-elements navbar-nav">
    
<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../../setup.html">
    Setup
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../../getting_started.html">
    Getting Started
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../../tutorials/index.html">
    Tutorials
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../../settings.html">
    Settings
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../../api_reference/index.html">
    API reference
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../../architecture.html">
    Architecture
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../../faq.html">
    FAQ
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../../about.html">
    About
  </a>
</li>

  </ul>
</nav></div>
      
    </div>
    
    
    <div class="navbar-header-items__end">
      
        <div class="navbar-item navbar-persistent--container">
          

 <script>
 document.write(`
   <button class="btn navbar-btn search-button-field search-button__button" title="Search" aria-label="Search" data-bs-placement="bottom" data-bs-toggle="tooltip">
    <i class="fa-solid fa-magnifying-glass"></i>
    <span class="search-button__default-text">Search</span>
    <span class="search-button__kbd-shortcut"><kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd class="kbd-shortcut__modifier">K</kbd></span>
   </button>
 `);
 </script>
        </div>
      
      
        <div class="navbar-item">

<script>
document.write(`
  <button class="btn btn-sm navbar-btn theme-switch-button" title="light/dark" aria-label="light/dark" data-bs-placement="bottom" data-bs-toggle="tooltip">
    <span class="theme-switch nav-link" data-mode="light"><i class="fa-solid fa-sun fa-lg"></i></span>
    <span class="theme-switch nav-link" data-mode="dark"><i class="fa-solid fa-moon fa-lg"></i></span>
    <span class="theme-switch nav-link" data-mode="auto"><i class="fa-solid fa-circle-half-stroke fa-lg"></i></span>
  </button>
`);
</script></div>
      
    </div>
    
  </div>
  
  
    <div class="navbar-persistent--mobile">

 <script>
 document.write(`
   <button class="btn navbar-btn search-button-field search-button__button" title="Search" aria-label="Search" data-bs-placement="bottom" data-bs-toggle="tooltip">
    <i class="fa-solid fa-magnifying-glass"></i>
    <span class="search-button__default-text">Search</span>
    <span class="search-button__kbd-shortcut"><kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd class="kbd-shortcut__modifier">K</kbd></span>
   </button>
 `);
 </script>
    </div>
  

  
</div>

    </header>
  

  <div class="bd-container">
    <div class="bd-container__inner bd-page-width">
      
      
      
        
      
      <div class="bd-sidebar-primary bd-sidebar hide-on-wide">
        

  
  <div class="sidebar-header-items sidebar-primary__section">
    
    
      <div class="sidebar-header-items__center">
        
          
          
            <div class="navbar-item">
<nav class="navbar-nav">
  <ul class="bd-navbar-elements navbar-nav">
    
<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../../setup.html">
    Setup
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../../getting_started.html">
    Getting Started
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../../tutorials/index.html">
    Tutorials
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../../settings.html">
    Settings
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../../api_reference/index.html">
    API reference
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../../architecture.html">
    Architecture
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../../faq.html">
    FAQ
  </a>
</li>


<li class="nav-item pst-header-nav-item">
  <a class="nav-link nav-internal" href="../../../../about.html">
    About
  </a>
</li>

  </ul>
</nav></div>
          
        
      </div>
    
    
    
      <div class="sidebar-header-items__end">
        
          <div class="navbar-item">

<script>
document.write(`
  <button class="btn btn-sm navbar-btn theme-switch-button" title="light/dark" aria-label="light/dark" data-bs-placement="bottom" data-bs-toggle="tooltip">
    <span class="theme-switch nav-link" data-mode="light"><i class="fa-solid fa-sun fa-lg"></i></span>
    <span class="theme-switch nav-link" data-mode="dark"><i class="fa-solid fa-moon fa-lg"></i></span>
    <span class="theme-switch nav-link" data-mode="auto"><i class="fa-solid fa-circle-half-stroke fa-lg"></i></span>
  </button>
`);
</script></div>
        
      </div>
    
  </div>
  
  
  <div class="sidebar-primary-items__end sidebar-primary__section">
  </div>
  
  <div id="rtd-footer-container"></div>


      </div>
      
      <main id="main-content" class="bd-main" role="main">
        
        
          <div class="bd-content">
            <div class="bd-article-container">
              
              <div class="bd-header-article d-print-none">
<div class="header-article-items header-article__inner">
  
    <div class="header-article-items__start">
      
        <div class="header-article-item">



<nav aria-label="Breadcrumb" class="d-print-none">
  <ul class="bd-breadcrumbs">
    
    <li class="breadcrumb-item breadcrumb-home">
      <a href="../../../../index.html" class="nav-link" aria-label="Home">
        <i class="fa-solid fa-home"></i>
      </a>
    </li>
    
    <li class="breadcrumb-item"><a href="../../../index.html" class="nav-link">Module code</a></li>
    
    <li class="breadcrumb-item active" aria-current="page">brevitas.cor...</li>
  </ul>
</nav>
</div>
      
    </div>
  
  
</div>
</div>
              
              
              
                
<div id="searchbox"></div>
                <article class="bd-article">
                  
  <h1>Source code for brevitas.core.bit_width.parameter</h1><div class="highlight"><pre>
<span></span><span class="c1"># Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.</span>
<span class="c1"># SPDX-License-Identifier: BSD-3-Clause</span>

<span class="kn">from</span><span class="w"> </span><span class="nn">typing</span><span class="w"> </span><span class="kn">import</span> <span class="n">Optional</span>

<span class="kn">import</span><span class="w"> </span><span class="nn">torch</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">torch</span><span class="w"> </span><span class="kn">import</span> <span class="n">Tensor</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">torch.nn</span><span class="w"> </span><span class="kn">import</span> <span class="n">Module</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">torch.nn</span><span class="w"> </span><span class="kn">import</span> <span class="n">Parameter</span>

<span class="kn">import</span><span class="w"> </span><span class="nn">brevitas</span>
<span class="kn">import</span><span class="w"> </span><span class="nn">brevitas.config</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">config</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">brevitas.core.function_wrapper</span><span class="w"> </span><span class="kn">import</span> <span class="n">RoundSte</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">brevitas.core.restrict_val</span><span class="w"> </span><span class="kn">import</span> <span class="n">IntRestrictValue</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">brevitas.function</span><span class="w"> </span><span class="kn">import</span> <span class="n">abs_binary_sign_grad</span>

<span class="n">MIN_INT_BIT_WIDTH</span> <span class="o">=</span> <span class="mi">2</span>
<span class="n">NON_ZERO_EPSILON</span> <span class="o">=</span> <span class="mf">1e-6</span>
<span class="n">REMOVE_ZERO_BIT_WIDTH</span> <span class="o">=</span> <span class="mf">0.1</span>


<div class="viewcode-block" id="BitWidthParameter"><a class="viewcode-back" href="../../../../api_reference/brevitas.core.bit_width.html#brevitas.core.bit_width.parameter.BitWidthParameter">[docs]</a><span class="k">class</span><span class="w"> </span><span class="nc">BitWidthParameter</span><span class="p">(</span><span class="n">brevitas</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">ScriptModule</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    ScriptModule that returns a learnable bit-width wrapped in a float torch.Tensor.</span>

<span class="sd">    Args:</span>
<span class="sd">        bit_width (int): value to initialize the output learned bit-width.</span>
<span class="sd">        min_bit_width (int): lower bound for the output learned bit-width. Default: 2.</span>
<span class="sd">        restrict_bit_width_impl: restrict the learned bit-width to a subset of values. Default: IntRestrictValue(RoundSte()).</span>
<span class="sd">        override_pretrained_bit_width (bool): ignore pretrained bit-width loaded from a state dict. Default: False.</span>

<span class="sd">    Returns:</span>
<span class="sd">        Tensor: bit-width wrapped in a float torch.tensor and backend by a learnable torch.nn.Parameter.</span>

<span class="sd">    Raises:</span>
<span class="sd">        RuntimeError: if bit_width &lt; min_bit_width.</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; bit_width_parameter = BitWidthParameter(8)</span>
<span class="sd">        &gt;&gt;&gt; bit_width_parameter()</span>
<span class="sd">        tensor(8., grad_fn=&lt;RoundSteFnBackward&gt;)</span>

<span class="sd">    Note:</span>
<span class="sd">        Set env variable BREVITAS_IGNORE_MISSING_KEYS=1 to avoid errors when retraining</span>
<span class="sd">        from a floating point state dict.</span>

<span class="sd">    Note:</span>
<span class="sd">        Maps to bit_width_impl_type == BitWidthImplType.PARAMETER == &#39;PARAMETER&#39; == &#39;parameter&#39; in higher-level APIs.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">__constants__</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;bit_width_base&#39;</span><span class="p">,</span> <span class="s1">&#39;override_pretrained&#39;</span><span class="p">]</span>

    <span class="k">def</span><span class="w"> </span><span class="fm">__init__</span><span class="p">(</span>
            <span class="bp">self</span><span class="p">,</span>
            <span class="n">bit_width</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
            <span class="n">min_bit_width</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="n">MIN_INT_BIT_WIDTH</span><span class="p">,</span>
            <span class="n">restrict_bit_width_impl</span><span class="p">:</span> <span class="n">Module</span> <span class="o">=</span> <span class="n">IntRestrictValue</span><span class="p">(</span><span class="n">RoundSte</span><span class="p">()),</span>
            <span class="n">override_pretrained_bit_width</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
            <span class="n">dtype</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">dtype</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
            <span class="n">device</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">device</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">BitWidthParameter</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

        <span class="k">if</span> <span class="n">bit_width</span> <span class="o">&lt;</span> <span class="n">MIN_INT_BIT_WIDTH</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span>
                <span class="s2">&quot;Int bit width has to be at least </span><span class="si">{}</span><span class="s2">, instead is </span><span class="si">{}</span><span class="s2">.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
                    <span class="n">MIN_INT_BIT_WIDTH</span><span class="p">,</span> <span class="n">bit_width</span><span class="p">))</span>

        <span class="k">if</span> <span class="n">min_bit_width</span> <span class="o">&lt;</span> <span class="n">MIN_INT_BIT_WIDTH</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span>
                <span class="s2">&quot;Min int bit width has to be at least </span><span class="si">{}</span><span class="s2">, instead is </span><span class="si">{}</span><span class="s2">.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
                    <span class="n">MIN_INT_BIT_WIDTH</span><span class="p">,</span> <span class="n">min_bit_width</span><span class="p">))</span>

        <span class="k">if</span> <span class="n">bit_width</span> <span class="o">&lt;</span> <span class="n">min_bit_width</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span>
                <span class="s2">&quot;Int bit width has to be at least </span><span class="si">{}</span><span class="s2">, instead is </span><span class="si">{}</span><span class="s2">.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
                    <span class="n">min_bit_width</span><span class="p">,</span> <span class="n">bit_width</span><span class="p">))</span>

        <span class="n">bit_width</span> <span class="o">=</span> <span class="nb">float</span><span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">bit_width</span><span class="p">))</span>
        <span class="n">min_bit_width</span> <span class="o">=</span> <span class="nb">float</span><span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">min_bit_width</span><span class="p">))</span>
        <span class="n">bit_width_base</span> <span class="o">=</span> <span class="n">restrict_bit_width_impl</span><span class="o">.</span><span class="n">restrict_init_float</span><span class="p">(</span><span class="n">min_bit_width</span><span class="p">)</span>
        <span class="n">bit_width</span> <span class="o">=</span> <span class="n">restrict_bit_width_impl</span><span class="o">.</span><span class="n">restrict_init_float</span><span class="p">(</span><span class="n">bit_width</span><span class="p">)</span>
        <span class="n">bit_width_offset_init</span> <span class="o">=</span> <span class="n">bit_width</span> <span class="o">-</span> <span class="n">bit_width_base</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">bit_width_offset</span> <span class="o">=</span> <span class="n">Parameter</span><span class="p">(</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">bit_width_offset_init</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">dtype</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">device</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">bit_width_base</span> <span class="o">=</span> <span class="n">bit_width_base</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">restrict_bit_width_impl</span> <span class="o">=</span> <span class="n">restrict_bit_width_impl</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">override_pretrained</span> <span class="o">=</span> <span class="n">override_pretrained_bit_width</span>

<div class="viewcode-block" id="BitWidthParameter.forward"><a class="viewcode-back" href="../../../../api_reference/brevitas.core.bit_width.html#brevitas.core.bit_width.parameter.BitWidthParameter.forward">[docs]</a>    <span class="nd">@brevitas</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">script_method</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="n">bit_width</span> <span class="o">=</span> <span class="n">abs_binary_sign_grad</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">bit_width_offset</span><span class="p">)</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">bit_width_base</span>
        <span class="n">bit_width</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">restrict_bit_width_impl</span><span class="p">(</span><span class="n">bit_width</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">bit_width</span></div>

    <span class="k">def</span><span class="w"> </span><span class="nf">_load_from_state_dict</span><span class="p">(</span>
            <span class="bp">self</span><span class="p">,</span> <span class="n">state_dict</span><span class="p">,</span> <span class="n">prefix</span><span class="p">,</span> <span class="n">local_metadata</span><span class="p">,</span> <span class="n">strict</span><span class="p">,</span> <span class="n">missing_keys</span><span class="p">,</span> <span class="n">unexpected_keys</span><span class="p">,</span>
            <span class="n">error_msgs</span><span class="p">):</span>
        <span class="n">bit_width_const_key</span> <span class="o">=</span> <span class="n">prefix</span> <span class="o">+</span> <span class="s1">&#39;bit_width&#39;</span>
        <span class="n">bit_width_offset_key</span> <span class="o">=</span> <span class="n">prefix</span> <span class="o">+</span> <span class="s1">&#39;bit_width_offset&#39;</span>
        <span class="k">if</span> <span class="n">bit_width_const_key</span> <span class="ow">in</span> <span class="n">state_dict</span><span class="p">:</span>
            <span class="k">assert</span> <span class="n">bit_width_offset_key</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">state_dict</span><span class="p">,</span> <span class="s2">&quot;Both should not be true.&quot;</span>
            <span class="n">bit_width</span> <span class="o">=</span> <span class="n">state_dict</span><span class="p">[</span><span class="n">bit_width_const_key</span><span class="p">]</span>
            <span class="n">state_dict</span><span class="p">[</span><span class="n">bit_width_offset_key</span><span class="p">]</span> <span class="o">=</span> <span class="n">bit_width</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">bit_width_base</span>
            <span class="k">del</span> <span class="n">state_dict</span><span class="p">[</span><span class="n">bit_width_const_key</span><span class="p">]</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">override_pretrained</span> <span class="ow">and</span> <span class="n">bit_width_offset_key</span> <span class="ow">in</span> <span class="n">state_dict</span><span class="p">:</span>
            <span class="k">del</span> <span class="n">state_dict</span><span class="p">[</span><span class="n">bit_width_offset_key</span><span class="p">]</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">BitWidthParameter</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="n">_load_from_state_dict</span><span class="p">(</span>
            <span class="n">state_dict</span><span class="p">,</span> <span class="n">prefix</span><span class="p">,</span> <span class="n">local_metadata</span><span class="p">,</span> <span class="n">strict</span><span class="p">,</span> <span class="n">missing_keys</span><span class="p">,</span> <span class="n">unexpected_keys</span><span class="p">,</span> <span class="n">error_msgs</span><span class="p">)</span>
        <span class="k">if</span> <span class="p">(</span><span class="n">config</span><span class="o">.</span><span class="n">IGNORE_MISSING_KEYS</span> <span class="ow">or</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">override_pretrained</span><span class="p">)</span> <span class="ow">and</span> <span class="n">bit_width_offset_key</span> <span class="ow">in</span> <span class="n">missing_keys</span><span class="p">:</span>
            <span class="n">missing_keys</span><span class="o">.</span><span class="n">remove</span><span class="p">(</span><span class="n">bit_width_offset_key</span><span class="p">)</span></div>


<div class="viewcode-block" id="RemoveBitwidthParameter"><a class="viewcode-back" href="../../../../api_reference/brevitas.core.bit_width.html#brevitas.core.bit_width.parameter.RemoveBitwidthParameter">[docs]</a><span class="k">class</span><span class="w"> </span><span class="nc">RemoveBitwidthParameter</span><span class="p">(</span><span class="n">brevitas</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">ScriptModule</span><span class="p">):</span>
    <span class="n">__constants__</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;non_zero_epsilon&#39;</span><span class="p">,</span> <span class="s1">&#39;override_pretrained&#39;</span><span class="p">]</span>

    <span class="k">def</span><span class="w"> </span><span class="fm">__init__</span><span class="p">(</span>
            <span class="bp">self</span><span class="p">,</span>
            <span class="n">bit_width_to_remove</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
            <span class="n">override_pretrained_bit_width</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
            <span class="n">non_zero_epsilon</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">NON_ZERO_EPSILON</span><span class="p">,</span>
            <span class="n">remove_zero_bit_width</span><span class="o">=</span><span class="n">REMOVE_ZERO_BIT_WIDTH</span><span class="p">,</span>
            <span class="n">dtype</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">dtype</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
            <span class="n">device</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">device</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">RemoveBitwidthParameter</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

        <span class="k">if</span> <span class="n">bit_width_to_remove</span> <span class="o">&lt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Bit width to clamp has to be &gt;= 0.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">bit_width_to_remove</span><span class="p">))</span>
        <span class="k">elif</span> <span class="n">bit_width_to_remove</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">bit_width_coeff_init</span> <span class="o">=</span> <span class="mi">1</span> <span class="o">/</span> <span class="n">remove_zero_bit_width</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">bit_width_coeff_init</span> <span class="o">=</span> <span class="mi">1</span> <span class="o">/</span> <span class="n">bit_width_to_remove</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">bit_width_coeff</span> <span class="o">=</span> <span class="n">Parameter</span><span class="p">(</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">bit_width_coeff_init</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">dtype</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">device</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">non_zero_epsilon</span> <span class="o">=</span> <span class="n">non_zero_epsilon</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">override_pretrained</span> <span class="o">=</span> <span class="n">override_pretrained_bit_width</span>

<div class="viewcode-block" id="RemoveBitwidthParameter.forward"><a class="viewcode-back" href="../../../../api_reference/brevitas.core.bit_width.html#brevitas.core.bit_width.parameter.RemoveBitwidthParameter.forward">[docs]</a>    <span class="nd">@brevitas</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">script_method</span>
    <span class="k">def</span><span class="w"> </span><span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
        <span class="n">bit_width_to_remove</span> <span class="o">=</span> <span class="mf">1.0</span> <span class="o">/</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">non_zero_epsilon</span> <span class="o">+</span> <span class="n">torch</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">bit_width_coeff</span><span class="p">))</span>
        <span class="k">return</span> <span class="n">bit_width_to_remove</span></div>

    <span class="k">def</span><span class="w"> </span><span class="nf">_load_from_state_dict</span><span class="p">(</span>
            <span class="bp">self</span><span class="p">,</span> <span class="n">state_dict</span><span class="p">,</span> <span class="n">prefix</span><span class="p">,</span> <span class="n">local_metadata</span><span class="p">,</span> <span class="n">strict</span><span class="p">,</span> <span class="n">missing_keys</span><span class="p">,</span> <span class="n">unexpected_keys</span><span class="p">,</span>
            <span class="n">error_msgs</span><span class="p">):</span>
        <span class="n">bit_width_coeff_key</span> <span class="o">=</span> <span class="n">prefix</span> <span class="o">+</span> <span class="s1">&#39;bit_width_coeff&#39;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">override_pretrained</span> <span class="ow">and</span> <span class="n">bit_width_coeff_key</span> <span class="ow">in</span> <span class="n">state_dict</span><span class="p">:</span>
            <span class="k">del</span> <span class="n">state_dict</span><span class="p">[</span><span class="n">bit_width_coeff_key</span><span class="p">]</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">RemoveBitwidthParameter</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="n">_load_from_state_dict</span><span class="p">(</span>
            <span class="n">state_dict</span><span class="p">,</span> <span class="n">prefix</span><span class="p">,</span> <span class="n">local_metadata</span><span class="p">,</span> <span class="n">strict</span><span class="p">,</span> <span class="n">missing_keys</span><span class="p">,</span> <span class="n">unexpected_keys</span><span class="p">,</span> <span class="n">error_msgs</span><span class="p">)</span>
        <span class="k">if</span> <span class="p">(</span><span class="n">config</span><span class="o">.</span><span class="n">IGNORE_MISSING_KEYS</span> <span class="ow">or</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">override_pretrained</span><span class="p">)</span> <span class="ow">and</span> <span class="n">bit_width_coeff_key</span> <span class="ow">in</span> <span class="n">missing_keys</span><span class="p">:</span>
            <span class="n">missing_keys</span><span class="o">.</span><span class="n">remove</span><span class="p">(</span><span class="n">bit_width_coeff_key</span><span class="p">)</span></div>
</pre></div>

                </article>
              
              
              
              
              
                <footer class="prev-next-footer d-print-none">
                  
<div class="prev-next-area">
</div>
                </footer>
              
            </div>
            
            
              
            
          </div>
          <footer class="bd-footer-content">
            
          </footer>
        
      </main>
    </div>
  </div>
  
  <!-- Scripts loaded after <body> so the DOM is not blocked -->
  <script src="../../../../_static/scripts/bootstrap.js?digest=3ee479438cf8b5e0d341"></script>
<script src="../../../../_static/scripts/pydata-sphinx-theme.js?digest=3ee479438cf8b5e0d341"></script>

  <footer class="bd-footer">
<div class="bd-footer__inner bd-page-width">
  
    <div class="footer-items__start">
      
        <div class="footer-item">

  <p class="copyright">
    
      © Copyright 2025 - Advanced Micro Devices, Inc..
      <br/>
    
  </p>
</div>
      
        <div class="footer-item">

  <p class="sphinx-version">
    Created using <a href="https://www.sphinx-doc.org/">Sphinx</a> 5.3.0.
    <br/>
  </p>
</div>
      
    </div>
  
  
  
    <div class="footer-items__end">
      
        <div class="footer-item">
<script>
document.write(`
  <div class="version-switcher__container dropdown">
    <button id="pst-version-switcher-button-2"
      type="button"
      class="version-switcher__button btn btn-sm navbar-btn dropdown-toggle"
      data-bs-toggle="dropdown"
      aria-haspopup="listbox"
      aria-controls="pst-version-switcher-list-2"
      aria-label="Version switcher list"
    >
      Choose version  <!-- this text may get changed later by javascript -->
      <span class="caret"></span>
    </button>
    <div id="pst-version-switcher-list-2"
      class="version-switcher__menu dropdown-menu list-group-flush py-0"
      role="listbox" aria-labelledby="pst-version-switcher-button-2">
      <!-- dropdown will be populated by javascript on page load -->
    </div>
  </div>
`);
</script></div>
      
    </div>
  
</div>

  </footer>
  </body>
</html>